# 基于IPR项目2的手写数字识别应用：Pinecone云服务版
# 功能：通过Gradio界面实现手写数字识别，使用Pinecone进行KNN查询(k=11)
# 前置要求：
# 1. 安装依赖: pip install pinecone gradio numpy pillow
# 2. 已创建Pinecone索引'mnist-index'并填充数据
# 3. 替换API_KEY为有效密钥

from __future__ import annotations
import io
import base64
from typing import Optional
import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone, ServerlessSpec

# Pinecone配置与连接管理
PINECONE_INDEX = "mnist-index"
API_KEY = "pcsk_4EcLgW_MpufGtAtj64LC3kASF1U834xnJP1M12k6ihSTyfkBteSHcGQpdTmjLVZzMkPRdb"

# 全局连接对象
pinecone_client: Optional[Pinecone] = None
index_connection = None


def setup_pinecone():
    """初始化Pinecone连接（单例模式）"""
    global pinecone_client, index_connection
    if index_connection is not None:
        return
    
    try:
        pinecone_client = Pinecone(api_key=API_KEY)
        index_connection = pinecone_client.Index(PINECONE_INDEX)
    except Exception as err:
        index_connection = None
        print(f"[错误] Pinecone连接失败: {str(err)}")
        print("提示: 请检查API密钥、网络连接和索引'mnist-index'是否存在")


# 图像预处理函数
def preprocess_image(img: Image.Image, threshold_val: float = 0.2, invert_force: bool = False) -> Optional[np.ndarray]:
    """将输入图像转换为8x8的特征数组（0-16范围）"""
    # 转为灰度图并归一化
    gray_img = img.convert('L')
    img_array = np.array(gray_img, dtype=np.float32)
    value_range = np.ptp(img_array)
    
    if value_range < 1e-8:
        return None
        
    normalized = (img_array - img_array.min()) / value_range
    inverted = 1.0 - normalized
    
    # 确定最佳前景/背景反转
    threshold_mask = normalized > threshold_val
    inverted_mask = inverted > threshold_val
    area_norm = threshold_mask.sum()
    area_inv = inverted_mask.sum()
    total_pixels = normalized.size

    def calculate_score(area: int) -> float:
        ratio = area / max(1, total_pixels)
        return -abs(ratio - 0.2)

    # 选择合适的图像版本
    if area_norm == 0 and area_inv == 0:
        return None
    if area_norm == 0:
        selected = inverted
    elif area_inv == 0:
        selected = normalized
    else:
        selected = normalized if calculate_score(area_norm) >= calculate_score(area_inv) else inverted

    if invert_force:
        selected = 1.0 - selected

    # 提取有效区域
    valid_mask = selected > threshold_val
    if not valid_mask.any():
        return None
        
    y_coords, x_coords = np.where(valid_mask)
    y_min, y_max = y_coords.min(), y_coords.max() + 1
    x_min, x_max = x_coords.min(), x_coords.max() + 1
    cropped = selected[y_min:y_max, x_min:x_max]

    # 缩放与居中
    h, w = cropped.shape
    max_side = max(h, w)
    target_size = 32
    scale_factor = (target_size - 4) / max_side
    new_h, new_w = max(1, round(h * scale_factor)), max(1, round(w * scale_factor))

    # 图像调整
    cropped_img = Image.fromarray((cropped * 255).astype(np.uint8))
    resampler = Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS
    resized = cropped_img.resize((new_w, new_h), resampler)

    # 放置到画布中央
    canvas = Image.new('L', (target_size, target_size), color=0)
    offset_x = (target_size - new_w) // 2
    offset_y = (target_size - new_h) // 2
    canvas.paste(resized, (offset_x, offset_y))

    # 转为8x8并映射到0-16
    small_img = canvas.resize((8, 8), resampler)
    small_array = np.array(small_img, dtype=np.float32)
    small_range = np.ptp(small_array) or 1.0
    final_array = (small_array - small_array.min()) / small_range * 16.0
    return final_array


# 推理函数
def recognize_digit(image: Optional[np.ndarray | Image.Image | dict], threshold: float = 0.2, invert: bool = False):
    """处理输入图像并通过Pinecone查询进行识别"""
    if image is None:
        return None, None

    # 初始化Pinecone连接
    setup_pinecone()
    if index_connection is None:
        return None, None

    # 处理不同类型的图像输入
    processed_img = None
    if isinstance(image, dict):
        # 处理Gradio Sketchpad的字典格式
        for key in ['image', 'composite', 'mask']:
            if key in image and image[key] is not None:
                processed_img = image[key]
                break
        if isinstance(processed_img, dict):
            for key in ['image', 'composite']:
                if key in processed_img and processed_img[key] is not None:
                    processed_img = processed_img[key]
                    break
        if isinstance(processed_img, str):
            # 处理base64编码图像
            try:
                if processed_img.startswith('data:image'):
                    _, b64_data = processed_img.split(',', 1)
                else:
                    b64_data = processed_img
                img_bytes = base64.b64decode(b64_data)
                processed_img = Image.open(io.BytesIO(img_bytes)).convert('RGBA')
            except Exception as e:
                print(f"[警告] 解码图像失败: {repr(e)}")
                return None, None
    elif isinstance(image, np.ndarray):
        processed_img = Image.fromarray(image.astype('uint8'))
    elif isinstance(image, Image.Image):
        processed_img = image
    else:
        print(f"[警告] 不支持的图像类型: {type(image)}")
        return None, None

    # 处理RGBA图像
    if processed_img.mode == 'RGBA':
        bg = Image.new('RGBA', processed_img.size, (255, 255, 255, 255))
        processed_img = Image.alpha_composite(bg, processed_img).convert('RGB')

    # 预处理为8x8特征
    feature_array = preprocess_image(processed_img, threshold_val=threshold, invert_force=invert)
    if feature_array is None:
        print("[警告] 预处理后图像为空，请输入更清晰的数字")
        return None, None

    print(f"[调试] 输入尺寸: {processed_img.size}, 特征范围: {feature_array.min():.2f}-{feature_array.max():.2f}, 均值: {feature_array.mean():.2f}, 和: {feature_array.sum():.2f}")

    # 准备查询向量
    query_vector = feature_array.flatten().astype(float).tolist()

    # 执行Pinecone查询
    try:
        query_result = index_connection.query(vector=query_vector, top_k=11, include_metadata=True)
        matches = query_result.get("matches", []) if isinstance(query_result, dict) else getattr(query_result, "matches", [])
        
        # 提取标签并投票
        label_list = []
        for match in matches:
            metadata = match.get("metadata") if isinstance(match, dict) else getattr(match, "metadata", {})
            if "Label" in metadata:
                label_list.append(int(metadata["Label"]))
        
        prediction = int(np.argmax(np.bincount(label_list))) if label_list else 0
    except Exception as e:
        print(f"[错误] 查询失败: {repr(e)}")
        return None, None

    # 生成可视化结果
    vis_array = (feature_array / 16.0 * 255.0).clip(0, 255).astype(np.uint8)
    vis_image = Image.fromarray(vis_array, mode='L').resize((128, 128), Image.NEAREST)
    return prediction, np.array(vis_image)


# Gradio界面
with gr.Blocks(title="手写数字识别（Pinecone KNN版）") as app:
    gr.Markdown("""
    # 手写数字识别（Pinecone KNN版）
    - 后端使用Pinecone云服务进行KNN搜索（k=11）
    - 首次使用可能较慢（需建立云端连接）
    - 请确保已准备好'mnist-index'索引及训练数据
    """)

    with gr.Row():
        drawing_pad = gr.Sketchpad(label="在此处手写数字(0-9)")
        result_display = gr.Number(label="识别结果", precision=0)
        preview_window = gr.Image(label="预处理效果", type="numpy")

    with gr.Row():
        threshold_control = gr.Slider(0.05, 0.5, 0.2, 0.05, label="前景阈值(值越小越敏感)")
        invert_checkbox = gr.Checkbox(label="强制反色", value=False)

    with gr.Row():
        submit_btn = gr.Button("识别")
        clear_btn = gr.Button("清除")

    submit_btn.click(recognize_digit, [drawing_pad, threshold_control, invert_checkbox], [result_display, preview_window])
    clear_btn.click(lambda: None, None, drawing_pad)

if __name__ == "__main__":
    app.launch()